import numpy as np
import sys
import random
import pandas as pd

def CutNoneNodule(predScan):
	newScan = dict()
	for scanName in predScan:
		scan = predScan[scanName]
		candidates = predScan[scanName]['candidates']
		scan['candidates'] = list()
		for cdd in candidates:
			if cdd['prob'] >= 0.10:
				scan['candidates'].append(cdd)

		newScan[scanName] = scan
	return newScan


def CheckConsist(predScanList):
	if len(predScanList) == 0:
		raise ValueError('Length must great than zero')
	baseScans = predScanList[0]['predScan']
	scanList = list(baseScans.keys())
	for scanName in scanList:
		# print(scanName)
		for i, predScan in enumerate(predScanList):
			if i == 0:
				continue
			
			scan = predScan['predScan'][scanName]
			baseScan = baseScans[scanName]
			assert scan['cancer'] == baseScan['cancer']
			# print(scanName)
			assert len(scan['candidates']) == len(baseScan['candidates'])
			for j, cdd in enumerate(baseScan['candidates']):
				# print(cdd)
				assert cdd['path'] == scan['candidates'][j]['path']


def GetScore(cddList):
	upperBound = 0.7
	lowerBound = 0.1
	score = lowerBound
	for cdd in cddList:
		if cdd['prob'] < 0.1:
			continue
		score = max(score, cdd['mal_prob'])
	score = score * 1.5
	score = min(score, upperBound)
	return score

def KaggleDivide(scanList):
	trainSet = []
	testSet = []
	for scanName in scanList:
		if scanList[scanName]['source'] == 'kaggle_train':
			trainSet.append(scanName)
		else:
			testSet.append(scanName)
	return trainSet, testSet


def loadScanLabel():
	trainLabel = 'stage1_labels.csv'
	testLabel = 'stage1_test_labels.csv'
	train_labcsv = pd.read_csv(trainLabel)
	test_labelcsv = pd.read_csv(testLabel)
	trainLabel = dict(train_labcsv.values)
	testLabel = dict(test_labelcsv.values)
	return trainLabel, testLabel

def cdds2scanStage2(cddsList):
	
	predScan = dict()
	kaggleTest = np.load('../../data/stage2/test_list.npy')
	for scanid in kaggleTest:
		predScan[scanid] = dict()
		predScan[scanid]['candidates'] = []
		predScan[scanid]['source'] = 'kaggle_test'

	for cdd in cddsList:
		scanid = cdd['scanid']
		# print(scanid)
		if cdd['prob'] <= 0.10:continue
		predScan[scanid]['candidates'].append(cdd)

	return predScan

def cdds2scan(cddsList, mode='train'):
	if mode == 'stage2':
		return cdds2scanStage2(cddsList)
	predScan = dict()
	kaggleTrain, kaggleTest = loadScanLabel()
	if mode == 'train' or mode == 'all':
		for scanid in kaggleTrain:
			predScan[scanid] = dict()
			predScan[scanid]['candidates'] = []
			predScan[scanid]['cancer'] = kaggleTrain[scanid]
			predScan[scanid]['source'] = 'kaggle_train'
	if mode == 'test' or mode == 'all':
		for scanid in kaggleTest:
			predScan[scanid] = dict()
			predScan[scanid]['candidates'] = []
			predScan[scanid]['cancer'] = kaggleTest[scanid]
			predScan[scanid]['source'] = 'kaggle_test'
	for cdd in cddsList:
		scanid = cdd['scanid']
		# print(scanid)
		if cdd['prob'] <= 0.10:continue
		predScan[scanid]['candidates'].append(cdd)
	return predScan

def DivideData(scanList, nFold=5, iFold=0):
	scanList = sorted(scanList)
	random.seed(0)
	random.shuffle(scanList)
	trainSet = []
	testSet = []

	for i, scanName in enumerate(scanList):
		if i%nFold == iFold:
			testSet.append(scanName)
		else:
			trainSet.append(scanName)
	return trainSet, testSet

def AvgOverScan(predScanList, scanid):
	avgScan = predScanList[0][scanid]
	for cdd in avgScan['candidates']:
		mean_p_mal = []
		for predScan in predScanList:
			for n_cdd in predScan[scanid]['candidates']:
				if n_cdd['path'] == cdd['path']:
					mean_p_mal.append(n_cdd['mal_prob'])
		assert len(mean_p_mal) == len(predScanList)
		cdd['mal_prob'] = np.mean(mean_p_mal)
	return avgScan



if __name__ == '__main__':
	loadScanLabel()